﻿#include <iostream>
#include "kalman_filter.h"
#include <Eigen/Dense>
#include "yolov8_detect.h"

int main()
{
    const int state_dim = 8;
    const int meas_dim = 4;
    std::string onnx_path = "F:/Code_cxx/tracking/resource/model/best-SiLU.onnx";
    std::unique_ptr<Yolov8Detect> hand_detect = std::make_unique<Yolov8Detect>();
    hand_detect->LoadModel("detect", onnx_path);

    KalmanFilter kf(state_dim, meas_dim);

    Eigen::VectorXd x0(state_dim);
    Eigen::MatrixXd P0 = Eigen::MatrixXd::Identity(state_dim, state_dim);
    Eigen::MatrixXd Q = Eigen::MatrixXd::Identity(state_dim, state_dim) * 0.1;
    Eigen::MatrixXd R = Eigen::MatrixXd::Identity(meas_dim, meas_dim) * 0.5;
    Eigen::MatrixXd F(state_dim, state_dim);

    x0 << 0, 0, 100, 100, 0, 0, 0, 0;

    F << 1, 0, 0, 0, 1, 0, 0, 0,
        0, 1, 0, 0, 0, 1, 0, 0,
        0, 0, 1, 0, 0, 0, 1, 0,
        0, 0, 0, 1, 0, 0, 0, 1,
        0, 0, 0, 0, 1, 0, 0, 0,
        0, 0, 0, 0, 0, 1, 0, 0,
        0, 0, 0, 0, 0, 0, 1, 0,
        0, 0, 0, 0, 0, 0, 0, 1;

    Eigen::MatrixXd H(meas_dim, state_dim);
    H << 1, 0, 0, 0, 0, 0, 0, 0,
        0, 1, 0, 0, 0, 0, 0, 0,
        0, 0, 1, 0, 0, 0, 0, 0,
        0, 0, 0, 1, 0, 0, 0, 0;

    // bool istracking = false;
    kf.init(x0, P0, Q, R, F, H);

    Eigen::VectorXd u = Eigen::VectorXd::Zero(state_dim);

    cv::Rect trackRect;

    cv::RNG rng(cv::getTickCount());

    cv::VideoCapture cap(0, cv::CAP_DSHOW);
    while (cap.isOpened())
    {
        cv::Mat frame;
        cap >> frame;
        cv::Mat letterboxMat;

        std::vector<int> padding;
        float scale = hand_detect->LetterBox(frame, letterboxMat, padding, cv::Size(640, 640), cv::Scalar(114, 114, 114), true);

        std::vector<DetectBox> bboxes = hand_detect->Inference(letterboxMat, scale, padding);

        for (const auto &box : bboxes)
        {
            if (box.IsValid())
            {
                cv::rectangle(frame, cv::Rect(box.left, box.top, box.right - box.left, box.bottom - box.top), cv::Scalar(0, 255, 0), 2);
                cv::putText(frame, box.label_name, cv::Point(box.left, box.top - 10), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 255, 0), 2);
                cv::putText(frame, std::to_string(box.score), cv::Point(box.left + 40, box.top - 10), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 255, 0), 2);

                trackRect = cv::Rect(box.left, box.top, box.right - box.left, box.bottom - box.top);

                // Initialize the Kalman filter with the center of the detection box
                // x0 << box.left + (box.right - box.left) / 2, box.top + (box.bottom - box.top) / 2,
                //     box.right - box.left, box.bottom - box.top, 0, 0, 0, 0;
                // kf.init(x0, P0);

                // 观察值
                Eigen::VectorXd z(meas_dim);
                // Use the center of the trackRect as measurement
                z << trackRect.x + trackRect.width / 2, trackRect.y + trackRect.height / 2, trackRect.width, trackRect.height;
                kf.predict(u);
                kf.update(z);
                Eigen::VectorXd state = kf.getState();

                // Update trackRect based on the new state
                cv::Rect updateRect = cv::Rect(state(0) - state(2) / 2, state(1) - state(3) / 2, state(2), state(3));
                std::cout << updateRect << std::endl;
                cv::rectangle(frame, updateRect, cv::Scalar(0, 0, 255), 2);
            }
        }

        if (cv::waitKey(1) == 27)
        {
            break;
        }
        cv::imshow("test", frame);
    }
    cap.release();
    cv::destroyAllWindows();

    return 0;
}